# -*- coding: utf-8 -*-
from typing import Optional, Any
import datetime
import numpy as np
from xarray import Dataset, DataArray
import bz2
import gzip

SWAN_HEADER = [
    ("data_type", "12c"),
    ("data_name", "38c"),
    ("name", "8c"),
    ("version", "8c"),
    ("year", "u2"),
    ("month", "u2"),
    ("day", "u2"),
    ("hour", "u2"),
    ("minute", "u2"),
    ("interval", "u2"),
    ("x_grid_num", "u2"),
    ("y_grid_num", "u2"),
    ("z_grid_num", "u2"),
    ("radar_count", "i4"),
    ("start_lon", "f4"),
    ("start_lat", "f4"),
    ("center_lon", "f4"),
    ("center_lat", "f4"),
    ("x_reso", "f4"),
    ("y_reso", "f4"),
    ("height", "40f4"),
    ("station_names", "20S16"),
    ("station_lons", "20f4"),
    ("station_lats", "20f4"),
    ("station_alts", "20f4"),
    ("mosaic_flags", "20B"),
    ("m_data_type", "i2"),
    ("dimension", "i2"),
    ("res", "168c"),
]

SWAN_dtype = np.dtype(SWAN_HEADER)


def prepare_file(file: Any) -> Any:
    if hasattr(file, "read"):
        return file
    f = open(file, "rb")
    magic = f.read(3)
    f.close()
    if magic.startswith(b"\x1f\x8b"):
        return gzip.GzipFile(file, "rb")
    if magic.startswith(b"BZh"):
        return bz2.BZ2File(file, "rb")
    return open(file, "rb")


class SWAN(object):
    r"""
    Class reading SWAN grid data.

    Args:
        file (str, IO): Path points to the file or a file object.
    """
    dtype_conv = {0: "B", 1: "b", 2: "u2", 3: "i2", 4: "u2"}
    size_conv = {0: 1, 1: 1, 2: 2, 3: 2, 4: 2}

    def __init__(self, file: Any, product: Optional[str] = None):
        f = prepare_file(file)
        header = np.frombuffer(f.read(1024), SWAN_dtype)
        xdim, ydim, zdim = (
            header["x_grid_num"][0],
            header["y_grid_num"][0],
            header["z_grid_num"][0],
        )
        dtype = header["m_data_type"][0]
        data_size = int(xdim) * int(ydim) * int(zdim) * self.size_conv[dtype]
        bittype = self.dtype_conv[dtype]
        data_body = np.frombuffer(f.read(data_size), bittype).astype(int)
        out = data_body.reshape(ydim, xdim)
        out[np.where(out == 1)] = 0

        self.data_time = datetime.datetime(
            header["year"].item(),
            header["month"].item(),
            header["day"].item(),
            header["hour"].item(),
            header["minute"].item(),
        )
        # TODO: Recognize correct product name
        self.product_name = (
            (b"".join(header["data_name"]).decode("gbk", "ignore").replace("\x00", ""))
            if not product
            else product
        )
        start_lon = header["start_lon"][0]
        start_lat = header["start_lat"][0]
        center_lon = header["center_lon"][0]
        center_lat = header["center_lat"][0]
        end_lon = center_lon * 2 - start_lon
        end_lat = center_lat * 2 - start_lat
        x_reso = header['x_reso'][0]
        y_reso = header['y_reso'][0]

        xstep = -x_reso if start_lon > end_lon else x_reso
        ystep = -y_reso if start_lat > end_lat else y_reso
        self.lon = np.arange(start_lon, end_lon, xstep)[:xdim]  # For shape compatibility
        self.lat = np.arange(start_lat, end_lat, ystep)[:ydim]
        self.data = (np.ma.masked_equal(out, 0) - 66) / 2

    def get_data(self) -> Dataset:
        r"""
        Get radar data with extra information

        Returns:
            xarray.Dataset: Data.
        """
        dtype = self.product_name
        ret = self.data

        da = DataArray(ret, coords=[self.lat, self.lon], dims=["latitude", "longitude"])
        ds = Dataset(
            {dtype: da},
            attrs={
                "scan_time": self.data_time.strftime("%Y-%m-%d %H:%M:%S"),
                "site_code": "SWAN",
                "site_name": "SWAN",
                "tangential_reso": np.nan,
                "range": np.nan,
                "elevation": 0,
            },
        )
        return ds

def get_rad(radfile):
    a2 = 118.239
    b2 = 1.5241
    min_z = -5.
    max_z = 55.

    f = SWAN(radfile)
    data = f.get_data()

    lat, lon, rad = data.variables.values()
    rad = np.minimum(max_z, rad)
    rad = np.maximum(min_z, rad)

    rad_r = np.power(10, (rad - 10 * np.log10(a2)) / (10 * b2))
    return np.nan_to_num(rad_r), lat.values, lon.values